import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
from monai import transforms
import json
import os
import clip
import SimpleITK as sitk


def datafold_read(datalist):
    with open(datalist) as f:
        json_data = json.load(f)

    return json_data

def convert_to_multi_channel(y):
    y_wt, y_tc, y_et = y > 0, ((y == 1) + (y == 3)) > 0, y == 3
    return np.concatenate([y_wt, y_tc, y_et], axis=0)
    

class Seg3DBraTs2023(Dataset):
    def __init__(self, data_list, phase='train'):
        super(Seg3DBraTs2023, self).__init__()
        self.data_list = data_list[phase]

        self.train_transform = transforms.Compose(
            [
                transforms.CropForegroundd(
                    keys=["image", "label"], source_key="image", k_divisible=[128, 128, 128]
                ),
                transforms.RandSpatialCropd(
                    keys=["image", "label"], roi_size=[128, 128, 128], random_size=False
                ),
                transforms.RandFlipd(keys=["image", "label"], prob=0.2, spatial_axis=0),
                transforms.RandFlipd(keys=["image", "label"], prob=0.2, spatial_axis=1),
                transforms.RandFlipd(keys=["image", "label"], prob=0.2, spatial_axis=2),
                transforms.NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
                transforms.RandScaleIntensityd(keys="image", factors=(-0.3, 0.3), prob=0.15),
                transforms.RandShiftIntensityd(keys="image", offsets=0.1, prob=0.1),
                transforms.RandGaussianNoised(keys="image", prob=0.15, mean=0.0, std=0.33),
                transforms.RandGaussianSmoothd(keys="image", sigma_x=(0.5, 1.5), sigma_y=(0.5, 1.5), sigma_z=(0.5, 1.5), prob=0.15),
                transforms.ToTensord(keys=["image", "label"]),
            ]
        )
        self.val_transform = transforms.Compose(
            [
                transforms.CropForegroundd(
                    keys=["image", "label"], source_key="image", k_divisible=[128, 128, 128]
                ),
                transforms.RandSpatialCropd(
                    keys=["image", "label"], roi_size=[128, 128, 128], random_size=False
                ),
                transforms.NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
                transforms.ToTensord(keys=["image", "label"]),
            ]
        )

        self.test_transform = transforms.Compose(
            [
                transforms.NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
                transforms.ToTensord(keys=["image", "label"]),
            ]
        )

        self.phase = phase

    def __len__(self):
        return len(self.data_list)

    def __getitem__(self, item):
        data = self.load_image(self.data_list[item])
        if self.phase == 'train':
            data = self.train_transform(data)
        elif self.phase == 'val':
            data = self.val_transform(data)
        elif self.phase == 'test':
            data = self.test_transform(data)

        return data

    def load_image(self, file_dic):
        image = []
        assert len(file_dic['image']) == 4, f"BraTS need four modality but get {len(file_dic['image'])}"

        for i in sorted(file_dic['image']):
            image_path = i
            image_ = sitk.ReadImage(image_path)
            image_ = sitk.GetArrayFromImage(image_)

            image.append(image_)
        image = np.stack(image, axis=0)

        label_path = file_dic['label']
        label = sitk.ReadImage(label_path)
        label = sitk.GetArrayFromImage(label)
        label = convert_to_multi_channel(label)

        return {
            'image': image,
            'label': label,
            'path': file_dic['image'][0]
        }
        

class Seg2DBraTs2023(Dataset):
    def __init__(self, data_list, phase='train'):
        super(Seg2DBraTs2023, self).__init__()
        self.data_list = data_list[phase]

        self.train_transform = transforms.Compose(
            [
                transforms.CropForegroundd(
                    keys=["image", "label"], source_key="image", k_divisible=[128, 128]
                ),
                transforms.RandSpatialCropd(
                    keys=["image", "label"], roi_size=[128, 128], random_size=False
                ),
                transforms.RandFlipd(keys=["image", "label"], prob=0.2, spatial_axis=0),
                transforms.RandFlipd(keys=["image", "label"], prob=0.2, spatial_axis=1),
                transforms.NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
                transforms.RandScaleIntensityd(keys="image", factors=(-0.3, 0.3), prob=0.15),
                transforms.RandShiftIntensityd(keys="image", offsets=0.1, prob=0.1),
                transforms.RandGaussianNoised(keys="image", prob=0.15, mean=0.0, std=0.33),
                transforms.RandGaussianSmoothd(keys="image", sigma_x=(0.5, 1.5), sigma_y=(0.5, 1.5), sigma_z=(0.5, 1.5), prob=0.15),
                transforms.ToTensord(keys=["image", "label"]),
            ]
        )
        self.val_transform = transforms.Compose(
            [
                transforms.CropForegroundd(
                    keys=["image", "label"], source_key="image", k_divisible=[128, 128]
                ),
                transforms.RandSpatialCropd(
                    keys=["image", "label"], roi_size=[128, 128], random_size=False
                ),
                transforms.NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
                transforms.ToTensord(keys=["image", "label"]),
            ]
        )

        self.test_transform = transforms.Compose(
            [
                transforms.NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
                transforms.ToTensord(keys=["image", "label"]),
            ]
        )

        self.phase = phase

    def __len__(self):
        return len(self.data_list)

    def __getitem__(self, item):
        data = self.load_image(self.data_list[item])
        if self.phase == 'train':
            data = self.train_transform(data)
        elif self.phase == 'val':
            data = self.val_transform(data)
        elif self.phase == 'test':
            data = self.test_transform(data)

        return data

    def load_image(self, file_dic):
        data = np.load(file_dic['path'])
        
        t1c = data['t1c']
        t1n = data['t1n']
        t2w = data['t2w']
        t2f = data['t2f']
        
        image = np.stack([t1c, t1n, t2w, t2f], axis=0).reshape(4, 256, 256)
        
        label = data['label']
        
        return {
            'image': image,
            'label': label,
            'path': os.path.split(file_dic['path'])[-1]
        }

class Imp2DBraTs2023(Dataset):
    def __init__(self, data_list, phase='train'):
        super(Imp2DBraTs2023, self).__init__()
        self.data_list = data_list[phase]

        self.train_transform = transforms.Compose(
            [
                transforms.NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
                transforms.RandScaleIntensityd(keys="image", factors=(-0.3, 0.3), prob=0.15),
                transforms.RandShiftIntensityd(keys="image", offsets=0.1, prob=0.1),
                transforms.RandGaussianNoised(keys="image", prob=0.15, mean=0.0, std=0.33),
                transforms.RandGaussianSmoothd(keys="image", sigma_x=(0.5, 1.5), sigma_y=(0.5, 1.5), sigma_z=(0.5, 1.5), prob=0.15),
                transforms.ToTensord(keys=["image", 'label']),
            ]
        )
        
        self.val_transform = transforms.Compose(
            [
                transforms.NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
                transforms.ToTensord(keys=["image", 'label']),
            ]
        )

        self.test_transform = transforms.Compose(
            [
                transforms.NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
                transforms.ToTensord(keys=["image", 'label']),
            ]
        )

        self.phase = phase

    def __len__(self):
        return len(self.data_list)

    def __getitem__(self, item):
        data = self.load_image(self.data_list[item])
        if self.phase == 'train':
            data = self.train_transform(data)
        elif self.phase == 'val':
            data = self.val_transform(data)
        elif self.phase == 'test':
            data = self.test_transform(data)

        return data

    def load_image(self, file_dic):
        data = np.load(file_dic['path'])
        
        t1c = data['t1c']
        t1n = data['t1n']
        t2w = data['t2w']
        t2f = data['t2f']
        
        image = np.stack([t1c, t1n, t2w, t2f], axis=0).reshape(4, 256, 256)
        
        label = data['label']
        
        if file_dic['wt'] is None:
            file_dic['wt'] = [0] * 5

        if file_dic['tc'] is None:
            file_dic['tc'] = [0] * 5
            
        if file_dic['et'] is None:
            file_dic['et'] = [0] * 5
        
        token = file_dic['wt'] + file_dic['tc'] + file_dic['et']
        token = clip.tokenize(str(token)).to(torch.float32).squeeze(0)

        return {
            'image': image,
            'label': label,
            'token': token,
            'path': os.path.split(file_dic['path'])[-1]
        }
               

class Gen2DBraTs2023(Dataset):
    def __init__(self, data_list, phase='train'):
        super(Gen2DBraTs2023, self).__init__()
        self.data_list = data_list[phase]

        self.train_transform = transforms.Compose(
            [
                transforms.NormalizeIntensityd(keys=['t1c', 't1n', 't2w', 't2f'], nonzero=True, channel_wise=True),
                transforms.RandScaleIntensityd(keys=['t1c', 't1n', 't2w', 't2f'], factors=(-0.3, 0.3), prob=0.15),
                transforms.RandShiftIntensityd(keys=['t1c', 't1n', 't2w', 't2f'], offsets=0.1, prob=0.1),
                transforms.RandGaussianNoised(keys=['t1c', 't1n', 't2w', 't2f'], prob=0.15, mean=0.0, std=0.33),
                transforms.RandGaussianSmoothd(keys=['t1c', 't1n', 't2w', 't2f'], sigma_x=(0.5, 1.5), sigma_y=(0.5, 1.5), sigma_z=(0.5, 1.5), prob=0.15),
                transforms.ToTensord(keys=['t1c', 't1n', 't2w', 't2f', 'label']),
            ]
        )
        
        self.val_transform = transforms.Compose(
            [
                transforms.NormalizeIntensityd(keys=['t1c', 't1n', 't2w', 't2f'], nonzero=True, channel_wise=True),
                transforms.ToTensord(keys=['t1c', 't1n', 't2w', 't2f', 'label']),
            ]
        )

        self.test_transform = transforms.Compose(
            [
                transforms.NormalizeIntensityd(keys=['t1c', 't1n', 't2w', 't2f'], nonzero=True, channel_wise=True),
                transforms.ToTensord(keys=['t1c', 't1n', 't2w', 't2f', 'label']),
            ]
        )

        self.phase = phase

    def __len__(self):
        return len(self.data_list)

    def __getitem__(self, item):
        data = self.load_image(self.data_list[item])
        if self.phase == 'train':
            data = self.train_transform(data)
        elif self.phase == 'val':
            data = self.val_transform(data)
        elif self.phase == 'test':
            data = self.test_transform(data)

        return data

    def load_image(self, file_dic):
        data = np.load(file_dic['path'])
        
        t1c = data['t1c']
        t1n = data['t1n']
        t2w = data['t2w']
        t2f = data['t2f']
        
        label = data['label']
        
        if file_dic['wt'] is None:
            file_dic['wt'] = [0] * 5

        if file_dic['tc'] is None:
            file_dic['tc'] = [0] * 5
            
        if file_dic['et'] is None:
            file_dic['et'] = [0] * 5
        
        token = file_dic['wt'] + file_dic['tc'] + file_dic['et']
        token = clip.tokenize(str(token)).to(torch.float32).squeeze(0)

        return {
            't1c': t1c,
            't1n': t1n,
            't2w': t2w,
            't2f': t2f,
            'label': label,
            'token': token,
            'path': os.path.split(file_dic['path'])[-1]
        }
        
def get_loader(datalist_json,
               batch_size,
               num_works,
               phase=None,
               datasets="Seg3DBraTs2023"):

    files = datafold_read(datalist=datalist_json)

    datasets = eval(datasets)(data_list=files, phase=phase)
    
    if phase != 'train':
        dataloader = DataLoader(datasets,
                                batch_size=batch_size,
                                num_workers=num_works,
                                pin_memory=True,
                                shuffle=False,
                                drop_last=True)
    else:
        dataloader = DataLoader(datasets,
                                batch_size=batch_size,
                                num_workers=num_works,
                                pin_memory=True,
                                shuffle=True,
                                drop_last=True)
    return dataloader
    
if __name__ == "__main__":
    train_file_path = '/home/qlc/raid/dataset/Brats2023/Adult_Glioma/modality_imputation.json'
    # val_file_path = '/home/qlc/raid/dataset/Brats2023/ag_2d/ag_t1_val.json'
    # test_file_path = '/home/qlc/raid/dataset/Brats2023/ag_2d/ag_t1_test.json'
    
    dataloader = get_loader(train_file_path, 4, 0, phase='train', datasets='Gen2DBraTs2023')
    import torch.nn as nn
    mse = nn.MSELoss()
    for i, data in enumerate(dataloader):
        img = data['t1c']
        
        loss = mse(img[0].unsqueeze(0), img[1].unsqueeze(0))
        b = torch.randn((1, 1, 256, 256))
        print(loss)
        loss = mse(img[0].unsqueeze(0), b)
        print(loss)
        print(img.size())
        assert 1 == 2

    